package cn.dataling.ai.component;

import jakarta.websocket.*;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import org.apache.commons.lang3.ObjectUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * websocket服务端
 */
@ServerEndpoint(value = "/websocket/{openId}")
@Component
public class WebSocketServer {


    private static ChatComponent chatComponent;

    @Autowired
    public void setChatComponent(ChatComponent chatComponent) {
        WebSocketServer.chatComponent = chatComponent;
    }

    /**
     * session缓存容器 key是sessionId value是 openId 对应的所有session
     */
    private static final Map<String, Session> SESSION_POOL = new ConcurrentHashMap<>();

    private static final Map<String, Boolean> FLUX_TAKE_WHILE_POOL = new ConcurrentHashMap<>();

    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketServer.class);

    /**
     * 消息发送
     *
     * @param message 消息
     * @param openId  客户端
     */
    public static void sendMessage(String message, String openId,String conversationId) {

        // 先停止之前的流（防止冲突）
        FLUX_TAKE_WHILE_POOL.put(openId, true);

        Flux<String> stringFlux = chatComponent.chatByStream(message, conversationId);

        stringFlux
                .doOnSubscribe(s -> {
                    sendMessage("[start]", openId);
                })
                .doOnCancel(() -> {
                    System.out.println("取消");
                })
                .takeWhile(s -> FLUX_TAKE_WHILE_POOL.getOrDefault(openId, true))
                .subscribe(
                        e -> {
                            sendMessage(e, openId);
                        },
                        err -> {
                            LOGGER.error(err.getMessage());
                        },
                        () -> {
                            sendMessage("[end]", openId);
                        });
    }

    private static void sendMessage(String message, String openId) {
        Session session = SESSION_POOL.get(openId);
        try {
            if (ObjectUtils.isNotEmpty(session) && session.isOpen()) {
                session.getBasicRemote().sendText(message);
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 新建连接时
     *
     * @param session 会话
     * @param openId  客户端
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("openId") String openId) {
        SESSION_POOL.put(openId, session);
        LOGGER.info("客户端 {} 会话ID {} 连接成功", openId, session.getId());
    }

    /**
     * 连接关闭时
     *
     * @param session 会话
     * @param openId  客户端
     */
    @OnClose
    public void onClose(Session session, @PathParam("openId") String openId) {
        SESSION_POOL.remove(openId);
        LOGGER.info("客户端 {} 会话ID {} 关闭成功", openId, session.getId());
    }

    /**
     * 收到消息时
     *
     * @param session 会话
     * @param message 消息
     * @param openId  客户端
     */
    @OnMessage
    public void onMessage(Session session, String message, @PathParam("openId") String openId) {
        if (message.equals("100")) {
            heart(session, openId);
        }
    }

    /**
     * 停止回答
     *
     * @param openId
     */
    public static Boolean stopAnswer(String openId) {
        return FLUX_TAKE_WHILE_POOL.put(openId, false);
    }

    /**
     * 发生异常是时
     *
     * @param session   会话
     * @param throwable 异常
     */
    @OnError
    public void onError(Session session, Throwable throwable) {
        SESSION_POOL.remove(session.getId());
        throwable.printStackTrace();
    }

    /**
     * 心跳
     *
     * @param session 会话
     * @param client  客户端
     */
    public void heart(Session session, String client) {
        try {
            session.getBasicRemote().sendText("200");
            LOGGER.info("给客户端 {} 会话ID {} 返回心跳成功", client, session.getId());
        } catch (IOException exception) {
            exception.printStackTrace();
        }
    }
}
